#!/usr/bin/env/python3
#coding=utf-8
#====================================================================
#   @Author:伍家俊
#   @TODO:  使用MDH法进行机器人正运动学推导与逆运动求解的脚本编写
#   @Reference: 《机器人学导论》原书第三版
#   @Data:  2021/7/7
#           2021/7/15   修改了__check_in_range()函数中的未对self.__Min_limits与self.__Max_limits中的数据转化为角度再用于比较的bug
#           2021/7/15   开放了rmatrix与Eular值之间相互转化的函数
#           2021/7/15   添加transform matrix与pose值之间相互转化的函数
#====================================================================
import numpy as np
import numpy.matlib
import math
#===================================================================================================
#   @Name: ArmKine
#   @Description: 这是一个关于(I-T-T-T-I)五轴操作臂正逆运动学的类,主要包含正运动学函数与逆运动学函数
#   @DH-factors:|    i   |   alpha_i-1    |   a_i-1    |   d_i    |     theta_i     |   offset_i    |
#               |    1   |      0         |     0      |    0     |     theta1      |       0       |
#               |    2   |      -pi/2     |     0      |    0     |     theta2      |     -pi/2     |
#               |    3   |      0         |     73.45  |    0     |     theta3      |       0       |
#               |    4   |      0         |     73.45  |    0     |     theta4      |     pi/2      |
#               |    5   |      pi/2      |     0      |    0     |     theta5      |       0       |
#   @Attention:此脚本只针对颜色小方块的抓取进行逆运动学推导，即第2、3、4关节会少一组解。
#===================================================================================================
PI = math.pi    
class ArmKine():
    '''
        求(I-T-T-T-I)五轴操作臂正逆运动学的类
        此类主要针对小方块抓取任务。
        #1 I模块的运动范围[-90° 90°]
    '''
    def __init__(self):
        self.__offset=[0,-PI/2,0,PI/2,0] #初始化偏置值
        self.__length=73.45
        self.__Min_limits=[-PI,-PI/2,-PI/2,-PI/2,-PI]
        self.__Max_limits=[+PI,+PI/2,+PI/2,+PI/2,PI]

    #========================================================
    #   @Function:fkine()
    #   @ToDo:  根据正运动学的推导公式完成正运动学的计算
    #   @Input: angles 长度为5的角度列表 角度制
    #   @Output: pose  长度为6的位姿列表 内容分别为[x,y,z,roll,pitch,yaw]
    #========================================================
    # ITTTI机械臂构型运动学正解
    def fkine(self,angles):
        angles_ =angles.copy()
        radians_=[]
        for (i,angle) in enumerate(angles_):
            radians_.append(math.radians(angle))
        theta1=radians_[0]+self.__offset[0]
        theta2=radians_[1]+self.__offset[1]
        theta3=radians_[2]+self.__offset[2]
        theta4=radians_[3]+self.__offset[3]
        theta5=radians_[4]+self.__offset[4]
        
        #计算需要用到的三角函数
        s1=math.sin(theta1);c1=math.cos(theta1);c2=math.cos(theta2);s2=math.sin(theta2);
        s5=math.sin(theta5);c5=math.cos(theta5);
        c234=math.cos(theta2+theta3+theta4);s234=math.sin(theta2+theta3+theta4);
        c23=math.cos(theta2+theta3);s23=math.sin(theta2+theta3);
        
        #计算p_x,p_y,p_z三个位置角度
        p_x=c1*(self.__length*c2+self.__length*c23);p_y=s1*(self.__length*c2+self.__length*c23);p_z=-self.__length*s23-self.__length*s2;
        
        #计算旋转矩阵中的元素以求解Z-Y-X欧拉角
        r31=-s234*c5; r11=c234*c1*c5-s1*s5; r21=s1*c234*c5+c1*s5;
        r32=s234*s5;r33=c234;r12=-s1*c5-c234*c1*s5;r22=c1*c5-c234+s1+s5;
        r13=c1*s234;r23=s1*s234;
        r_list=[r11,r21,r31,r32,r33,r12,r22]
        Eular=self.__rmatrix2Eular(r_list)
        #pose=[p_x,p_y,p_z,Eular[0],Eular[1],Eular[2]]
        [p_x,p_y,p_z]=self.__base_and_tool_coordinate_transform([p_x,p_y,p_z],[r13,r23,r33])
        pose = [p_x,p_y,p_z,Eular[0],Eular[1],Eular[2]]
        return pose


    # ITTTI机械臂构型运动学逆解
    def ikine(self,poses):
        p_x=poses[0];p_y=poses[1];p_z=poses[2];
        roll=poses[3];pitch=poses[4];yaw=poses[5];
        rmatrix=numpy.matlib.zeros((3,3))
        Eular=[roll,pitch,yaw]
        #=============================
        #   rmatrix=[[r11,r12,r13],
        #            [r21,r22,r23],
        #            [r31,r32,r33]]
        #=============================
        rmatrix=self.__Eular2rmatrix(Eular)
        [p_x,p_y,p_z]=self.__base_and_tool_coordinate_transform_inv([p_x,p_y,p_z],[rmatrix[0,2],rmatrix[1,2],rmatrix[2,2]])
        
        #=============================================
        #   @ToDo:求解theta1
        #   @Attention:针对抓取任务仅需[-90° 90°]即可
        #=============================================
        theta1=math.atan2(p_y,p_x)      #theta1 为rad
        s1=math.sin(theta1);c1=math.cos(theta1);
        r11=rmatrix[0,0];r21=rmatrix[1,0];r12=rmatrix[0,1];r22=rmatrix[1,1];

        #求解theta5
        s5=-s1*r11+c1*r21;c5=-s1*r12+c1*r22;
        theta5=math.atan2(s5,c5);
        #求解theta2 theta3 theta4
        #为简化求解现在选择忽略theta2 theta3 theta4 的另外一组解
        #   见《机器人学导论》图4-10
        r13=rmatrix[0,2];r23=rmatrix[1,2];r33=rmatrix[2,2];
        c234=r33;s234=c1*r13+s1*r23;
        theta234=math.atan2(s234,c234);
        #求解theta3
        c3=(p_x**2+p_y**2+p_z**2-2*self.__length**2)/(2*self.__length**2)
        c3=round(c3,2)
        #当c3算得大于1.0和小于-1.0时认为运动学无解
        #if ((c3>1.0)or(c3<-1.0)):
        #    angle=[90,90,0,90,0]
        #    return angle
        theta3=math.atan2(math.sqrt(1-c3**2),c3)
        s3=math.sin(theta3)
        #求解theta2
        theta2=-math.atan2(p_z,math.sqrt(p_x**2+p_y**2))-math.atan2(self.__length*s3,self.__length+self.__length*c3)
        
        #求解theta4
        theta4=theta234-theta2-theta3
        
        angle=[]
        theta=[theta1,theta2,theta3,theta4,theta5]
        for i in range(5):
            angle.append(math.degrees(theta[i]-self.__offset[i]))
        #=====================================================
        #   若结果为0则赋值为0
        #=====================================================
        for i in range(5):
            if (self.__equal_or_not(angle[i],0,0.01)):
                angle[i]=0
            else:
                continue
        #=====================================================
        #   根据各关节范围来确定最终的计算结果
        #=====================================================
        for i in range(5):
            flag = self.__check_in_range(angle[i],i)
            if flag:
                continue
            else:
                if angle[i]>2*PI:
                    angle[i]-=2*PI
                if angle[i]<-2*PI:
                    angle[i]+=2*PI
        return angle

    #==========================================================
    #   @Function:__equal_or_not(self,angle,target.error)
    #   @ToDo:  检查某个值是否等于目标角度值
    #   @Input: angle   被测的角度
    #           target  目标角度
    #   @Output:True or False 相等为True,不等为False
    #==========================================================
    def __equal_or_not(self,angle,target,error):
        if (angle<=(target+error)) and (angle>=(target-error)): #即 检查angle值是否等于target值
            flag=True
        else:
            flag=False
        return flag

    #==========================================================
    #   @Function:__rmatrix2Eular(self,r_list):
    #   @ToDo:  根据特定的旋转矩阵元素值求解Z-Y-X欧拉角
    #   @Input: r_list  rotation_matrix_elements_lists 旋转矩阵中r11,r21,r31,r32,r33,r12,r22所组成的列表
    #   @Output: Eular 列表 [roll,pitch,yaw]
    #==========================================================
    def __rmatrix2Eular(self,r_list):
        '''
        r_list=[r11,r21,r31,r32,r33,r12,r22]
        '''
        pitch = math.degrees(math.atan2(-r_list[2],math.sqrt(r_list[0]**2+r_list[1]**2)))
        if self.__equal_or_not(pitch,90,0.01):
            roll=0.0
            yaw=math.degrees(math.atan2(r_list[5],r_list[6]))
        elif self.__equal_or_not(pitch,-90,0.01):
            roll=0.0
            yaw =-math.degrees(math.atan2(r_list[5],r_list[6]))
        else:
            roll = math.degrees(math.atan2(r_list[1],r_list[0]))
            yaw =math.degrees(math.atan2(r_list[3],r_list[4]))
        return [roll,pitch,yaw]
    
    #==========================================================
    #   @Function:__Eular2rmatrix():
    #   @ToDo    : 实现将Z-Y-X欧拉角转换为旋转矩阵
    #   @Input   : Eular 列表 [roll,pitch,yaw]
    #   @Output  : rmatrix 3x3矩阵
    #              [[r11,r12,r13],
    #               [r21,r22,r23],
    #               [r31,r32,r33]]
    #==========================================================
    def __Eular2rmatrix(self,Eular):
        
        #获取roll，pitch，yaw计算旋转矩阵
        rmatrix=numpy.matlib.zeros((3,3))
        roll=math.radians(Eular[0]);pitch=math.radians(Eular[1]);yaw=math.radians(Eular[2]);
        #计算aphla,beta,gama的cos与sin值 详见机器人学导论P32
        c_aphla=math.cos(roll);s_aphla=math.sin(roll);
        c_beta=math.cos(pitch);s_beta=math.sin(pitch);
        c_gama=math.cos(yaw);s_gama=math.sin(yaw);
        #计算rmatrix
        rmatrix[0,0]=c_aphla*c_beta;rmatrix[0,1]=c_aphla*s_beta*s_gama-s_aphla*c_gama;rmatrix[0,2]=c_aphla*s_beta*c_gama+s_aphla*s_gama;
        rmatrix[1,0]=s_aphla*c_beta;rmatrix[1,1]=s_aphla*s_beta*s_gama+c_aphla*c_gama;rmatrix[1,2]=s_aphla*s_beta*c_gama-c_aphla*s_gama;
        rmatrix[2,0]=-s_beta;rmatrix[2,1]=c_beta*s_gama;rmatrix[2,2]=c_beta*c_gama;
        for i in range(3):
            for j in range(3):
               if (self.__equal_or_not(rmatrix[i,j],0,0.0001)):
                   rmatrix[i,j]=0
               else:
                   continue
        return rmatrix

    #========================================================
    #   @Function:__check_in_range(self,angle,n)
    #   @Input:angle    输入的被检查的角度
    #   @Output:Flag    在范围中返回true,不在范围中返回false
    #========================================================
    def __check_in_range(self,angle,n):
        if (angle<=math.degrees(self.__Max_limits[n])-0.01) and (angle>=math.degrees(self.__Min_limits[n])+0.01): #即 检查是否符合 min_angle<=angle<=max_angle
            flag=True
        else:
            flag=False
        return flag

    #========================================================================================================
    #   @Function:__base_and_tool_coordinate_transform(self,[x,y,z],[r13,r23,r33])
    #   @Input: [x,y,z] 存储正运动学计算得到的位置坐标
    #           [r13,r23,r33] 存储正运动学计算得到的r13,r23,r33
    #   @Output: pose_res 列表 存储末端执行器在世界坐标系中的位置
    #   @Description: fkine计算的是机器人坐标系模型中{5}坐标系相对于与{1}坐标系原点重合的{0}坐标系的位置与姿态；
    #                 因此需要对fkine的结果进行处理使得{0}坐标系移动到机器人底座中,以及使得{5}坐标系变换到{tool}中
    #==========================================================================================================
    def __base_and_tool_coordinate_transform(self,argv1,argv2):
        '''
            argv1 List [x,y,z]
            argv2 List [r13,r23,r33]
        '''
        [p_x,p_y,p_z]=argv1.copy()
        [r13,r23,r33]=argv2.copy()
        p_x_tool2base = p_x+153.93*r13
        p_y_tool2base = p_y+153.93*r23
        p_z_tool2base = p_z+153.93*r33+224.7
        p_res = [p_x_tool2base,p_y_tool2base,p_z_tool2base]
        return p_res

    #==================================================================================================================
    #   @Function:__base_and_tool_coordinate_transform_inv(self,[p_x,p_y,p_z],[r13,r23,r33])
    #   @Input: [p_x,p_y,p_z] 存储逆运动学输入的位置值
    #           [r13,r23,r33] 存储逆运动学输入中旋转矩阵的r13,r23,r33值
    #   @Ouput: [p_x_input,p_y_input,p_z_input] 列表 准备传入ikine进行求解的位置值
    #   @Description: ikine计算的是已知机器人坐标系模型中{5}坐标系相对于原点与{1}坐标系原点重合的{0}坐标系的位置与姿态；
    #                 因此需要对ikine的输入进行处理使{tool}坐标系变换到{5}坐标系中，而{base}坐标系变换到{1}坐标系中
    #   @Call by: ikine()
    #==================================================================================================================
    def __base_and_tool_coordinate_transform_inv(self,argv1,argv2):
        '''
           argv1 list [p_x,p_y,p_z]
           argv2 list [r13,r23,r33]
        '''
        [p_x,p_y,p_z]=argv1.copy()
        [r13,r23,r33]=argv2.copy()
        p_x_input = -r13*153.93+p_x
        p_y_input = -r23*153.93+p_y
        p_z_input = -r33*153.93+p_z-224.7
        p_res = [p_x_input,p_y_input,p_z_input]
        return p_res
    

    def Eular2rmatrix(self,Eular):
        '''
       @Description: Eular2ramtrix函数的作用是将欧拉角转换为旋转矩阵
       @Argument:   Eular 列表 [roll,pitch,yaw] 分别代表Z-Y-X欧拉角中的三个角
       @Return:     rmarix 3x3矩阵 
       @Call method: 
            Robot=ArmKine()
            Eular=[roll,pitch,yaw]
            rmatrix=Robot.Eular2rmatrix(Eular)
        '''
        #获取roll，pitch，yaw计算旋转矩阵
        rmatrix=numpy.matlib.zeros((3,3))
        roll=math.radians(Eular[0]);pitch=math.radians(Eular[1]);yaw=math.radians(Eular[2]);
        #计算aphla,beta,gama的cos与sin值 详见机器人学导论P32
        c_aphla=math.cos(roll);s_aphla=math.sin(roll);
        c_beta=math.cos(pitch);s_beta=math.sin(pitch);
        c_gama=math.cos(yaw);s_gama=math.sin(yaw);
        #计算rmatrix
        rmatrix[0,0]=c_aphla*c_beta;rmatrix[0,1]=c_aphla*s_beta*s_gama-s_aphla*c_gama;rmatrix[0,2]=c_aphla*s_beta*c_gama+s_aphla*s_gama;
        rmatrix[1,0]=s_aphla*c_beta;rmatrix[1,1]=s_aphla*s_beta*s_gama+c_aphla*c_gama;rmatrix[1,2]=s_aphla*s_beta*c_gama-c_aphla*s_gama;
        rmatrix[2,0]=-s_beta;rmatrix[2,1]=c_beta*s_gama;rmatrix[2,2]=c_beta*c_gama;
        for i in range(3):
            for j in range(3):
               if (self.__equal_or_not(rmatrix[i,j],0,0.0001)):
                   rmatrix[i,j]=0
               else:
                   continue
        return rmatrix


    def rmatrix2Eular(self,r_list):
        '''
        @Description: rmatrix2Eular函数的作用是将旋转矩阵转化为欧拉角
        @Argument:  r_list list [r11,r21,r31,r32,r33,r12,r22]  仅存放旋转矩阵中个别值
        @Return:    Eular list [roll,pitch,yaw] Z-Y-X欧拉角
        @Call method:
            Robot=ArmKine()
            r_list=[rmatrix[0,0],rmatrix[1,0],rmatrix[2,0],rmatrix[2,1],rmatrix[2,2],rmatrix[0,1],rmatrix[1,1]]
            Eular=Robot.rmatrix2Eular(r_list)
        '''
        pitch = math.degrees(math.atan2(-r_list[2],math.sqrt(r_list[0]**2+r_list[1]**2)))
        if self.__equal_or_not(pitch,90,0.01):
            roll=0.0
            yaw=math.degrees(math.atan2(r_list[5],r_list[6]))
        elif self.__equal_or_not(pitch,-90,0.01):
            roll=0.0
            yaw=-math.degrees(math.atan2(r_list[5],r_list[6]))
        else:
            roll= math.degrees(math.atan2(r_list[1],r_list[0]))
            yaw =math.degrees(math.atan2(r_list[3],r_list[4]))
        return [roll,pitch,yaw]

    def tmatrix2pose(self,tmatrix):
        '''
        @Description: tmatrix2pose()实现将一个齐次变换矩阵转换为一个pose列表
        @Argument: tmatrix numpy.matlib.zeros((4,4)) [[r11,r12,r13,p_x],[r21,r22,r23,p_y],[r31,r32,r33,p_z],[0,0,0,1]]
        @Return:   pose list [x,y,z,roll,pitch,yaw]
        @Call method:
            Robot=ArmKine()
            pose=tmatrix2pose(tmatrix)
        '''
        pose=[0,0,0]
        rmatrix=np.matlib.zeros((3,3))
        for i in range(3):
            for j in range(3):
                rmatrix[i,j]=tmatrix[i,j]
        r_list=[rmatrix[0,0],rmatrix[1,0],rmatrix[2,0],rmatrix[2,1],rmatrix[2,2],rmatrix[0,1],rmatrix[1,1]]
        eular=self.__rmatrix2Eular(r_list)
        for i in range(3):
            pose[i]=tmatrix[i,3]
        for i in range(3):
            pose.append(eular[i])
        return pose

    def pose2tmatrix(self,pose):
        '''
        @Description: pose2tmatrix()实现将一个pose列表转换为一个齐次变换矩阵
        @Argument:   pose list [x,y,z,roll,pitch,yaw]
        @Return: tmatrix numpy.matlib.zeros((4,4)) [[r11,r12,r13,p_x],[r21,r22,r23,p_y],[r31,r32,r33,p_z],[0,0,0,1]]
        @Call method:
            Robot=ArmKine()
            tmatrix=pose2tmatrix(pose)
        '''
        tmatrix=np.matlib.zeros((4,4))
        for i in range(3):
            tmatrix[i,3]=pose[i]
        eular=[pose[3],pose[4],pose[5]]
        rmatrix=self.__Eular2rmatrix(eular)
        for i in range(3):
            for j in range(3):
                tmatrix[i,j]=rmatrix[i,j]
        tmatrix[3,3]=1;
        return tmatrix

